/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.fn.harness;



import com.google.auto.service.AutoService;

import com.bff.gaia.unified.fn.harness.control.BundleSplitListener;

import com.bff.gaia.unified.fn.harness.data.UnifiedFnDataClient;

import com.bff.gaia.unified.fn.harness.data.PCollectionConsumerRegistry;

import com.bff.gaia.unified.fn.harness.data.PTransformFunctionRegistry;

import com.bff.gaia.unified.fn.harness.state.UnifiedFnStateClient;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.Coder;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.PCollection;

import com.bff.gaia.unified.runners.core.construction.PTransformTranslation;

import com.bff.gaia.unified.sdk.fn.data.FnDataReceiver;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;



import java.io.IOException;

import java.util.Map;

import java.util.function.Supplier;



import static com.bff.gaia.unified.vendor.guava.com.google.common.collect.Iterables.getOnlyElement;



/** Executes flatten PTransforms. */

public class FlattenRunner<InputT> {

  /** A registrar which provides a factory to handle flatten PTransforms. */

  @AutoService(PTransformRunnerFactory.Registrar.class)

  public static class Registrar implements PTransformRunnerFactory.Registrar {



    @Override

    public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {

      return ImmutableMap.of(PTransformTranslation.FLATTEN_TRANSFORM_URN, new Factory());

    }

  }



  /** A factory for {@link FlattenRunner}. */

  static class Factory<InputT> implements PTransformRunnerFactory<FlattenRunner<InputT>> {

    @Override

    public FlattenRunner<InputT> createRunnerForPTransform(

        PipelineOptions pipelineOptions,

        UnifiedFnDataClient unifiedFnDataClient,

        UnifiedFnStateClient unifiedFnStateClient,

        String pTransformId,

        RunnerApi.PTransform pTransform,

        Supplier<String> processBundleInstructionId,

        Map<String, PCollection> pCollections,

        Map<String, Coder> coders,

        Map<String, RunnerApi.WindowingStrategy> windowingStrategies,

        PCollectionConsumerRegistry pCollectionConsumerRegistry,

        PTransformFunctionRegistry startFunctionRegistry,

        PTransformFunctionRegistry finishFunctionRegistry,

        BundleSplitListener splitListener)

        throws IOException {



      // Give each input a MultiplexingFnDataReceiver to all outputs of the flatten.

      String output = getOnlyElement(pTransform.getOutputsMap().values());

      FnDataReceiver<WindowedValue<?>> receiver =

          pCollectionConsumerRegistry.getMultiplexingConsumer(output);



      FlattenRunner<InputT> runner = new FlattenRunner<>();



      for (String pCollectionId : pTransform.getInputsMap().values()) {

        pCollectionConsumerRegistry.register(

            pCollectionId, pTransformId, (FnDataReceiver) receiver);

      }



      return runner;

    }

  }

}